// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#include <melon/benchmark.h>
#include <melon/init/init.h>

#include <pollux/functions/registerer.h>
#include <pollux/functions/lib/benchmarks/FunctionBenchmarkBase.h>
#include <pollux/functions/prestosql/registration/registration_functions.h>
#include <pollux/vector/vector_fuzzer.h>
#include <pollux/testing/vector/vector_test_base.h>

using namespace kumo::pollux;
using namespace kumo::pollux::exec;

namespace {

// Assumes flat arrays with flat elements placed sequentially with no gaps. Also
// assumes arrays on the left and on the rights sides have matching sizes and
// offsets. Evaluates on all rows.
VectorPtr evaluateFast(const RowVector& data) {
  auto left = data.childAt(0)->as_unchecked<ArrayVector>();
  auto right = data.childAt(1)->as_unchecked<ArrayVector>();
  auto numElements = left->elements()->size();
  auto rawLeft =
      left->elements()->as_unchecked<FlatVector<int64_t>>()->rawValues();
  auto rawRight =
      right->elements()->as_unchecked<FlatVector<int64_t>>()->rawValues();

  auto result = BaseVector::create(BIGINT(), numElements, data.pool());
  auto flatResult = result->as_unchecked<FlatVector<int64_t>>();
  auto rawResults = flatResult->mutableRawValues();

  for (auto i = 0; i < numElements; ++i) {
    rawResults[i] = rawLeft[i] + rawRight[i];
  }

  return std::make_shared<ArrayVector>(
      left->pool(),
      left->type(),
      nullptr,
      left->size(),
      left->offsets(),
      left->sizes(),
      result);
}

class ZipWithBenchmark : public functions::test::FunctionBenchmarkBase {
 public:
  explicit ZipWithBenchmark(uint32_t seed)
      : FunctionBenchmarkBase(), seed_{seed} {
    functions::prestosql::registerAllScalarFunctions();
  }

  RowVectorPtr generateData() {
    VectorFuzzer::Options options;
    options.vectorSize = 10'024;

    VectorFuzzer fuzzer(options, pool(), seed_);

    return vectorMaker_.row_vector(
        {fuzzer.fuzzFlat(ARRAY(BIGINT())), fuzzer.fuzzFlat(ARRAY(BIGINT()))});
  }

  void test() {
    auto data = generateData();

    auto basicResult = evaluate(kBasicExpression, data);
    auto fastResult = evaluateFast(*data);

    test::assertEqualVectors(basicResult, fastResult);
  }

  size_t runBasic(size_t times) {
    melon::BenchmarkSuspender suspender;
    auto data = generateData();
    auto exprSet = compileExpression(kBasicExpression, as_row_type(data->type()));
    suspender.dismiss();

    return doRun(exprSet, data, times);
  }

  size_t runFast(size_t times) {
    melon::BenchmarkSuspender suspender;
    auto data = generateData();
    suspender.dismiss();

    int cnt = 0;
    for (auto i = 0; i < times * 1'000; i++) {
      cnt += evaluateFast(*data)->size();
    }
    return cnt;
  }

 private:
  static const std::string kBasicExpression;

  size_t doRun(ExprSet& exprSet, const RowVectorPtr& row_vector, size_t times) {
    int cnt = 0;
    for (auto i = 0; i < times * 1'000; i++) {
      cnt += evaluate(exprSet, row_vector)->size();
    }
    return cnt;
  }

  const uint32_t seed_;
};

const std::string ZipWithBenchmark::kBasicExpression =
    "zip_with(c0, c1, (x, y) -> x + y)";

const uint32_t seed = melon::Random::rand32();

BENCHMARK_MULTI(basic, n) {
  ZipWithBenchmark benchmark(seed);
  return benchmark.runBasic(n);
}

BENCHMARK_MULTI(fast, n) {
  ZipWithBenchmark benchmark(seed);
  return benchmark.runFast(n);
}

} // namespace

int main(int argc, char** argv) {
  melon::Init init{&argc, &argv};

  KLOG(ERROR) << "Seed: " << seed;
  {
    ZipWithBenchmark benchmark(seed);
    benchmark.test();
  }
  melon::runBenchmarks();
  return 0;
}
